import numpy
import image_lib
import sys
import pylab



size_min_y=8
size_min_x=2
digit_size=128
max_height=1000

def l1_norm(i):
    d=sum(abs(i))
    while numpy.size(d) > 1:
        d=sum(d)
    return d 

def l2_norm(i):
    """ calculate l2 norm of i """
    d=sum(i**2)
    while numpy.size(d) > 1:
        d=sum(d)
    return d 

def get_neg(i):
    i2=1-i
    return i2    


def make_quadratic(i, qlen):
    """ transforms an image/rectangle to quadrat"""
    xlen=i.shape[1]
    ylen=i.shape[0]
    if (ylen > qlen):
        ylen=qlen
    if (xlen > qlen):
        xlen=qlen
    #print i.shape
    i2=numpy.zeros([qlen, qlen])
    for x in xrange(0,xlen):
        for y in xrange(0, ylen):
            i2[y,x]=i[y,x]
    return i2        

def remove_border(i):
    """ set border of image to zero """
    xlen=i.shape[1]
    ylen=i.shape[0]
    i[0:ylen, 0]=0
    i[0:ylen, xlen-1]=0
    i[0, 0:xlen]=0
    i[ylen-1, 0:xlen]=0
    return i
             
def locate_cell(i):
    #    flood_y=0
    #    x_len=i.shape[1]
    #    y_len=i.shape[0]
    #    y_max=int(y_len/4)
    #    ok=1
    #    y=0
    #    while ok:
    #        if i[y,x_len//2]>=1:
    #            ok=2
    #            flood_y=y+1
    #        
    #        if (ok==2):
    #            if not(i[y,x_len//2]>=1):
    #                ok=0
    #        if (y > y_len//2):
    #            ok=0
    #        y+=1
    #    i2=image_lib.flood_fill(i, x_len//2, flood_y, 0, 0.5)
    #    return i2
    #i*=0.5
    i=(1.0-i)
    i2=image_lib.segment_picture(i, 1.0)
    #bestimme objekt mit groesster flaeche- meisten pixel
    hi=numpy.histogram(i2, i2.max()+1)
    color_index=hi[0].argmax()
    #print hi
    i3=(i2==color_index)
    i3=i3.astype(float)
    #pylab.figure()
    #pylab.matshow(i3)
    return i3

def extract_number(i):
    min_width=4
    min_height=10
    i2=image_lib.segment_picture(i, 1.0)

    hi=image_lib.histogramm(i2)
    maxh=0
    idx=0
    #pylab.figure()
    #pylab.plot(hi)
    #pylab.matshow(i2)
    i4=None
    for n in xrange(1, len(hi)):
        if n > 0:
            i3=(i2==n)
            i3=i3.astype(float)
            #i3/=float(i3.max())
            #pylab.matshow(i3)
            minx,miny,maxx,maxy=image_lib.get_koordinates_roi(i3)

            w=abs(maxx-minx)
            d=abs(maxy-miny)
            #object is not connected to the border
            if (minx >0) and (maxx < i2.shape[1]-1):
                if (miny > 0) and (maxy < i2.shape[0]-1):
                    if w > min_width:
                        if d > min_height:
                            i4=i3.copy()
                            return i4
                    
    #print hi


    #pylab.matshow(i3)
    return i4
   

def compensate_rotation(i):
    i2=image_lib.reduce_size_roi(i)
    #print i2
    #pylab.figure()
    #pylab.imshow(i2)
    #print i
    x_len=i2.shape[1]
    h1=0
    x=int(x_len/3.0)
    while i2[h1, x]>=1.0:
        h1+=1
    x=int(x_len*2/3)
    h2=0
    while i2[h2, x]>=1.0:
        h2+=1
    h=h2-h1
    alpha=numpy.arctan(h/(x_len/3.0))
    #print alpha, h1, h2
    i3=numpy.misc.imrotate(i2, alpha)
    i3=(i3>0)
    i3=i3.astype(float)
    return i3
    
def close_image(i):
    i2=i.copy()
    s=i.shape
    for y in xrange(1, s[0]-1):
        for x in xrange(1, s[1]-1):

            ns=sum(sum(i[y-1:y+1,x-1:x+1]))
            
            if ns >= 3:
                i2[y,x]=1.0
            else:
                i2[y,x]=0.0
    return i2
                
def border_fill(i):
    x_len=i.shape[1]
    y_len=i.shape[0]
    for y in xrange(0, y_len):
        i[y,0:3]=1.0
        i[y, x_len-4:x_len]=1.0
    for x in xrange(0, x_len):
        i[0:3, x]=1.0
        i[y_len-4:y_len, x]=1.0
    return i

def calc_row_sum(i):
    y_len=i.shape[0]
    i2=numpy.zeros(y_len)
    for y in xrange(0, y_len):
        i2[y]=sum(i[y,:])
    return i2

def calc_column_sum(i):
    xlen=i.shape[1]
    i2=numpy.zeros([xlen])
    for x in xrange(xlen):
        i2[x]=sum(i[:,x])
    return i2

def get_skeleton(i):
    i2=i.copy()
    ok=1
    while ok:
        i3=i2.copy()
        ok=0
        for x in xrange(1, i2.shape[1]-1):
            for y in xrange(1, i2.shape[0]-1):
                s=sum(sum(i2[y-1:y+2,x-1:x+2]))
                #s=sum(i2[y-1:y+2,x])+sum(i2[y,x-1:x+2])-i2[y,x]
                #print s, i2[y-1:y+2,x-1:x+2]
                if  (s > 3) and (s < 8):
                    if i2[y,x]>=1:
                        i2[y,x]=0
                        ok=1
        #i2=i3
        #pylab.matshow(i2)
        #print i2
    return i2

def avg_filter(s):
    s2=numpy.zeros(len(s))
    s2[0]=s[0]
    for n in xrange(1,len(s)):
        s2[n]=(s[n-1]+s[n])/2.0
    return s2


def get_feature_vec(i):
    """feature vector: row and column sum """
    i2=calc_row_sum(i)
    i2=avg_filter(i2)
    i3=calc_column_sum(i)
    i3=avg_filter(i3)
    #print i2,"+",i3
    i4=numpy.concatenate((i3,i2))
    nf=float(i4.max())
    i4/=nf
    return i4

def get_feature_vec2(i):
    """ feature sum of pixels in image and sum of quaters"""
    s=i.shape
    v0=l2_norm(i)
    v1=l2_norm(i[0:s[0]/2, 0:s[1]/2])
    v2=l2_norm(i[0:s[0]/2, s[1]/2:s[1]])
    v3=l2_norm(i[s[0]/2:s[0], s[1]/2:s[1]])
    v4=l2_norm(i[s[0]/2:s[0], 0:s[1]/2])
    fv=numpy.array([v0, v1, v2, v3, v4])
    mv=float(fv.max())
    fv/=mv
    return fv

def get_feature_vec3(i):
    """ image is divided in 3x3 sub areas, sum of sub-area is calculated"""
    s=i.shape
    lx=s[1]/3
    ly=s[0]/3
    fv=numpy.zeros(3*3)
    for n in range(3):
        for m in range(3):
            v=i[n*ly:(n+1)*ly, m*lx:(m+1)*lx].sum()
            fv[m+n*3]=v
    fv/=float(fv.max())
    #print fv
    return fv
    
def get_feature_vec4(i):
    """image is divided in several areas, sum and diff is generated see viola&jones"""
    s=i.shape
    m=s[1]/2
    fv=numpy.zeros(20)
    fv[0]=i[:,0:m].sum()
    fv[1]=i[:,m:].sum()
    fv[2]=fv[1]-fv[2]
    m=s[0]/2
    fv[3]=i[0:m,:].sum()
    fv[4]=i[m:,:].sum()
    fv[5]=fv[3]-fv[4]
    m=s[1]/3
    fv[5]=i[:,0:m].sum()
    fv[6]=i[:,m:2*m].sum()
    fv[7]=i[:,2*m:].sum()
    fv[8]=fv[5]+fv[7]-fv[6]
    fv[9]=fv[5]-fv[7]-fv[6]
    m=s[0]/3
    fv[10]=i[0:m,:].sum()
    fv[11]=i[m:2*m,:].sum()
    fv[12]=i[2*m:,:].sum()
    fv[13]=fv[10]+fv[12]-fv[11]
    fv[14]=fv[10]-fv[12]-fv[11]
    m1=s[0]/2
    m2=s[1]/2
    fv[15]=i[:m1, :m2].sum()
    fv[16]=i[m1:,:m2].sum()
    fv[17]=i[:m1, m2:].sum()
    fv[18]=i[m1, m2:].sum()
    fv[19]=fv[15]-fv[16]+fv[17]-fv[18]
    fv/=float(fv.max())
    #print fv
    return fv

def get_filled_cell(i):
    i2=image_lib.flood_fill(i, 0,0, 0, 2)
    i2=(i2==2)
    i2=i2.astype(float)
    i2=get_neg(i2)
    return i2

def invert_image(i):
    i2=(1-i.copy())
    return i2

class ocr():
    def __init__(self, parent=None):
        self.parent=parent
        self.init_number_dict()
        
    def init_number_dict(self): 
        #pylab.ion()
        #i=image_lib.jaehne_test_pattern(400)
        #i=numpy.misc.imread("img_6373.jpg")
        i=image_lib.load("numbers.jpg")
        #print i
        #i=numpy.lena().copy()
        #i=numpy.misc.imresize(i, 0.2)
        i=image_lib.convert_to_grayscale(i)
        
        #pylab.gray()
        #hi=numpy.histogram(i, 256)
        #fb_index=numpy.argmax(hi[0])
        #fb_value=hi[1][fb_index]
        #i2=image_lib.binarize(i, fb_value)
        fb_value=i.min()+(i.max()-i.min())/2.0
        i2=(i<fb_value)
        i2=i2.astype(float)
        #pylab.imshow(i)
        i3=i2#i3=image_lib.flood_fill(i2, 0, 0, 0, 0.5)
        i4=image_lib.segment_picture(i3, 1)
        #print i4.min(), i4.max()
        #ylab.matshow(i4)
        #pylab.figure()
        #print hi
        #pylab.plot(hi[0])
        number_dict=dict()
        #pylab.figure()
        #pylab.imshow(i4)
        maxv=int(i4.max())
        #pylab.figure()
        #print maxv
        for n in range(1, maxv):
            #print  n
            i5=(i4==(n+1))
            #i5=i4
            i5=i5.astype(float)
            minx,miny,maxx, maxy=image_lib.get_koordinates_roi(i5)
            i6=i5[miny-1:maxy+2, minx-1:maxx+2].copy()
            #print minx,miny,maxx,maxy
            if (i6.shape[0] >2) and (i6.shape[1] > 2) :
        
                #i6=image_lib.flood_fill(i6, 0,0, 0, 0.5)
                #pylab.matshow(i6)
                #i6=(i6==0.5)
                #i6=i6.astype(float)
                #i6=(1-i6)
                #pylab.matshow(i6)
                #i6=image_lib.skeleton_pic(i6)
                #pylab.matshow(i6)
                #i6=image_lib.skeleton_pic(i6)
                i6=get_filled_cell(i6)
                i6=image_lib.reduce_size_roi(i6)
                scale_fac=digit_size/float(i6.shape[0])
                i6=image_lib.resize(i6, [int(scale_fac*i6.shape[1]), digit_size])
        
                #pylab.subplot(1, maxv, n)
                #i7=image_lib.calc_row_sum_horizontal(i6)
                #pylab.matshow(i6)
                #i7=make_quadratic(i6.copy(), digit_size)
                
                #pylab.matshow(i7)
                #i7=calc_row_sum(i7)
                #i7/=float(i7.max())
                #i7=get_neg(i7)
                #pylab.matshow(i7)
                #i6=get_feature_vec3(i6)
                i7=get_feature_vec4(i6)
                #print i7,i6
                i8=i7#numpy.concatenate((i6,i7))
                #pylab.figure()
                #pylab.plot(i7)
                #print i7.shape
                #print n
                number_dict[n]=i8.copy()#image_lib.get_entropy(i7)
                fnr=910+(n)
                #print fnr
                #pylab.subplot(fnr)
                #pylab.matshow(i7)
                #print number_dict[n-1]
                #pylab.plot(i7)
                #e=image_lib.get_histo_entropy(i6)
                #print "entropy ", e, " ", n
        #print number_dict
        #print "fertig!"
        #pylab.show()
        #print number_dict.keys()
        self.number_dict=number_dict
        return
    
    def get_sudoku_from_img(self, filename=""):    
        #pylab.ion()
        if filename == "":
            s_mat=numpy.zeros([9,9])
        else:
            if self.parent!=None:
                self.parent.generate_progress_dialog(100)
            #number_dict=init_number_dict()
            #print "init"  
            #i=image_lib.jaehne_test_pattern(400)
            #i=numpy.misc.imread("img_6373.jpg")
            i=image_lib.load(filename)
    
            if i.shape[0] > max_height:
                scale=(int(float(max_height)/i.shape[0]*i.shape[1]), max_height)
                #print i.shape, scale
                i=image_lib.resize(i, scale)
                #pylab.ion()
                #pylab.imshow(i)
                #print i
            #i=numpy.lena().copy()
            #i=numpy.misc.imresize(i, 0.2)
            #print i
            i=image_lib.convert_to_grayscale(i)
            #print "grayscale"
            i=image_lib.white_balance(i)
            #print "whitebalance"
            #grenzwert fuer vorder- hintergrund bestimmen
            fb_value=i.min()+(i.max()-i.min())/2.0
            #voder hintergrund trennen
            i2=(i<fb_value)
            i2=i2.astype(float)
            #hi2=numpy.histogram(i2, 256)
            #i2=remove_border(i2)
            #print "empty border"
            #i2=close_image(i2)
            #print "close"
            #print i2.shape
            i2=image_lib.reduce_size_roi(i2)
            #print i2.shape
            #pylab.figure()
            #pylab.matshow(i2)
            #i2=invert_image(i2)
            #i2=compensate_rotation(i2)
            ##pylab.figure()
            #
            #i2=invert_image(i2)
            #i2=close_image(i2)
            
            #print i2.shape
            #i2=image_lib.reduce_size_roi(i2)
            x_len=i2.shape[1]/9.0
            y_len=i2.shape[0]/9.0
            #pylab.ion()
            #print i2.shape
            #pylab.gray()
            #pylab.matshow(i2)
            #gitter raus
            pc=0
            s_mat=numpy.zeros([9,9])
            for y in xrange(0, 9):
                for x in xrange(0,9):
                    i3=i2[y*y_len:(y+1)*y_len, x*x_len:(x+1)*x_len].copy()
                    #pylab.matshow(i3)
                    #pylab.figure()
                    #pylab.imshow(i3)
                    if self.parent!=None:
                        self.parent.inc_progress()
                    #i4=locate_cell(i3)
                    #i5=(i4==1.0)
                    #i5=i5.astype(float)
                    #i5=image_lib.reduce_size_roi(i5)
                    #i5=invert_image(i5)
                    #pylab.matshow(i5)
                    i5=extract_number(i3)
                    #if i5 != None:
                    #    pylab.matshow(i5)
                    s_mat[y,x]=0
                    if i5 != None:
                        #print i5.shape
                        #pylab.matshow(i5)
                        #print xmin
                        #i5=image_lib.skeleton_pic(i5)
                        #pylab.matshow(i5)
                        #i5=image_lib.skeleton_pic(i5)
                        i5=get_filled_cell(i5)
                        i5=image_lib.reduce_size_roi(i5)
                        #pylab.matshow(i5)
                        if (i5.shape[0]>size_min_y) and (i5.shape[1]>size_min_x):
                            
                            scale_fac=digit_size/float(i5.shape[0])
                            i6=image_lib.resize(i5, [int(scale_fac*i5.shape[1]), digit_size])
                            #i6=numpy.ones([32,32])
                            #pylab.figure()
                            #pylab.matshow(i6)
                            #i7=image_lib.calc_row_sum_horizontal(i6)
                            #i7=make_quadratic(i6.copy(), digit_size)
                            #i7=calc_row_sum(i7)
                            #i7=numpy.ones([32,32])
                            #i7/=float(i7.max())
                            #pylab.matshow(i7)
                            #i7=get_neg(i7)
                            #i6=get_feature_vec3(i6)
                            i7=get_feature_vec4(i6)
                            i8=i7#numpy.concatenate((i6,i7))
                            #pylab.figure()
                            #pylab.plot(i7)
                            #print i7.shape
                            #m=image_lib.get_entropy(i7)
                           # m1=i7.mean()
                            d=numpy.inf
                            number=0
                            #print i7
                            #print i7
                            #i7=get_feature_vec(i7)
                            #pylab.figure()
                            #pylab.plot(i7)
                            #m1=i7.mean()
                            for n in self.number_dict.keys():
                                #print i7.shape, number_dict[n].shape
                                #i8=i7#self.number_dict[n]*i7
                                #m2=i8.mean()
                                #dn=l2_norm((i7-number_dict[n]))
                                #m2=i7.mean()#number_dict[n].mean()
                                #i8=i7*number_dict[n]
                                #pylab.matshow(i8)
    
                                dn=l2_norm((i8-self.number_dict[n]))
    
                                #dn=(l2_norm(i7-number_dict[n])+l2_norm(i8-number_dict[n]))/2
                                #print dn
                                #print i8.shape
                                #print dn
                                #dn=d
                
                                if dn < d:
                                    #pylab.matshow(i8)
                                    #i9=i8.copy()
                                    number=n
                                    d=dn
                            s_mat[y,x]=number
                            #i8=i7*number_dict[number]
                            #pylab.matshow(i8)
                            #pylab.figure()
                            #pylab.plot(i7)
                            #print "entropy ", image_lib.get_entropy(i7)
            
            if self.parent!=None:
                self.parent.close_progress_dialog()
            #print flood_x
            #pylab.show()
        return s_mat

fname=["2012-07-02sudoku_sz.jpg", "sudoku1.jpg","sudoku2.jpg","sudoku3.jpg","sudoku4.jpg", "sudoku5.jpg"]
  
if __name__=="__main__":
    pylab.gray()
    socr=ocr()
    if len(sys.argv) > 1:
        print sys.argv[1]
        
        #init_number_dict()
        #mat=1
        mat=socr.get_sudoku_from_img(sys.argv[1])
        #a=numpy.zeros([32,32])
        #a[10:25, 0:16]=1
        #pylab.matshow(a)
        #a=get_skeleton(a)
        #pylab.matshow(a)
        print mat 
    else:
        #pylab.gray()
        for n in fname:
            print n
            mat=socr.get_sudoku_from_img(n)
            print mat 
        
       
    pylab.show()
